#include <assert.h>
#include <algorithm>
#include <iomanip>
#include "champsim.h"
#include "cache.h"
#include "memory_class.h"
#include "feature_state.h"
#include "atropos.h"
#include "util.h"

#if 0
#	define LOCKED(...) {fflush(stdout); __VA_ARGS__; fflush(stdout);}
#	define LOGID() fprintf(stdout, "[%25s@%3u] ", \
							__FUNCTION__, __LINE__ \
							);
#	define MYLOG(...) LOCKED(LOGID(); fprintf(stdout, __VA_ARGS__); fprintf(stdout, "\n");)
#else
#	define MYLOG(...) {}
#endif

namespace knob
{
	extern float    atropos_alpha;
	extern float    atropos_gamma;
	extern float    atropos_epsilon;
	extern uint32_t atropos_state_num_bits;
	extern uint32_t atropos_max_states;
	extern uint32_t atropos_seed;
	extern string   atropos_policy;
	extern string   atropos_learning_type;
	extern uint32_t atropos_pt_size;
	extern uint32_t atropos_st_size;
	extern uint32_t atropos_ht_size;
	extern uint32_t atropos_max_pcs;
	extern uint32_t atropos_max_offsets;
	extern uint32_t atropos_max_deltas;
    extern uint32_t atropos_max_actions;
	extern int32_t  atropos_reward_none;
	extern int32_t  atropos_reward_incorrect;
	extern int32_t  atropos_reward_correct_untimely;
	extern int32_t  atropos_reward_correct_timely;
	extern int32_t  atropos_reward_expert;
	extern int32_t  atropos_latency_bias;
	extern int32_t  atropos_lookback_window;
	extern bool		atropos_brain_zero_init;
	extern bool     atropos_enable_reward_all;
	extern bool     atropos_enable_track_multiple;
	extern bool     atropos_enable_reward_out_of_bounds;
	extern int32_t  atropos_reward_out_of_bounds;
	extern uint32_t atropos_state_type;
	extern bool     atropos_access_debug;
	extern bool     atropos_print_access_debug;
	extern uint64_t atropos_print_access_debug_pc;
	extern uint32_t atropos_print_access_debug_pc_count;
	extern bool     atropos_print_trace;
	extern bool     atropos_enable_state_action_stats;
	extern bool     atropos_enable_reward_tracker_hit;
	extern int32_t  atropos_reward_tracker_hit;
	extern uint32_t atropos_state_hash_type;
	extern bool     atropos_enable_featurewise_engine;
	extern bool     atropos_enable_hierarchical_engine;
	extern uint32_t atropos_pref_degree;
	extern bool     atropos_enable_dyn_degree;
	extern vector<float> atropos_max_to_avg_q_thresholds;
	extern vector<int32_t> atropos_dyn_degrees;
	extern uint64_t atropos_early_exploration_window;
	extern uint32_t atropos_pt_address_hash_type;
	extern uint32_t atropos_pt_address_hashed_bits;
	extern uint32_t atropos_bloom_filter_size;
	extern uint32_t atropos_multi_deg_select_type;
	extern vector<int32_t> atropos_last_pref_offset_conf_thresholds;
	extern vector<int32_t> atropos_dyn_degrees_type2;
	extern uint32_t atropos_action_tracker_size;
	extern uint32_t atropos_high_bw_thresh;
	extern bool     atropos_enable_hbw_reward;
	extern int32_t  atropos_reward_hbw_correct_timely;
	extern int32_t  atropos_reward_hbw_correct_untimely;
	extern int32_t  atropos_reward_hbw_incorrect;
	extern int32_t  atropos_reward_hbw_none;
	extern int32_t  atropos_reward_hbw_out_of_bounds;
	extern int32_t  atropos_reward_hbw_tracker_hit;
	extern vector<int32_t> atropos_last_pref_offset_conf_thresholds_hbw;
	extern vector<int32_t> atropos_dyn_degrees_type2_hbw;

	/* Learning Engine knobs */
	extern bool     le_enable_trace;
	extern uint32_t le_trace_interval;
	extern std::string   le_trace_file_name;
	extern uint32_t le_trace_state;
	extern bool     le_enable_score_plot;
	extern std::vector<int32_t> le_plot_actions;
	extern std::string   le_plot_file_name;
	extern bool     le_enable_action_trace;
	extern uint32_t le_action_trace_interval;
	extern std::string le_action_trace_name;
	extern bool     le_enable_action_plot;

	/* Featurewise Engine knobs */
	extern vector<int32_t> 	le_featurewise_active_features;
	extern vector<int32_t> 	le_featurewise_active_worker;
	extern vector<int32_t> 	le_featurewise_num_tilings;
	extern vector<int32_t> 	le_featurewise_num_tiles;
	extern vector<int32_t> 	le_featurewise_hash_types;
	extern vector<int32_t> 	le_featurewise_enable_tiling_offset;
	extern float			le_featurewise_max_q_thresh;
	extern bool				le_featurewise_enable_action_fallback;
	extern vector<float> 	le_featurewise_feature_weights;
	extern bool				le_featurewise_enable_dynamic_weight;
	extern float			le_featurewise_weight_gradient;
	extern bool				le_featurewise_disable_adjust_weight_all_features_align;
	extern bool				le_featurewise_selective_update;
	extern uint32_t 		le_featurewise_pooling_type;
	extern bool             le_featurewise_enable_dyn_action_fallback;
	extern uint32_t 		le_featurewise_bw_acc_check_level;
	extern uint32_t 		le_featurewise_acc_thresh;
	extern bool 			le_featurewise_enable_trace;
	extern uint32_t		le_featurewise_trace_feature_type;
	extern string 			le_featurewise_trace_feature;
	extern uint32_t 		le_featurewise_trace_interval;
	extern uint32_t 		le_featurewise_trace_record_count;
	extern std::string 	le_featurewise_trace_file_name;
	extern bool 			le_featurewise_enable_score_plot;
	extern vector<int32_t> le_featurewise_plot_actions;
	extern std::string 	le_featurewise_plot_file_name;
	extern bool 			le_featurewise_remove_plot_script;

    extern uint8_t          le_page_shifting;
    extern uint8_t          le_offset_shifting;
}

void Atropos::init_knobs()
{
	if(knob::atropos_access_debug)
	{
		cout << "***WARNING*** setting knob::atropos_max_pcs, knob::atropos_max_offsets, and knob::atropos_max_deltas to large value as knob::atropos_access_debug is true" << endl;
		knob::atropos_max_pcs = 1024;
		knob::atropos_max_offsets = 1024;
		knob::atropos_max_deltas = 1024;
	}
	assert(knob::atropos_pref_degree >= 1 && (knob::atropos_pref_degree == 1 || !knob::atropos_enable_dyn_degree));
	assert(knob::atropos_max_to_avg_q_thresholds.size() == knob::atropos_dyn_degrees.size()-1);
	assert(knob::atropos_last_pref_offset_conf_thresholds.size() == knob::atropos_dyn_degrees_type2.size()-1);
}

void Atropos::init_stats()
{
	bzero(&stats, sizeof(stats));
	state_action_dist.clear();
}

Atropos::Atropos(string type, CACHE *cache) : Prefetcher(type), m_parent_cache(cache)
{
	init_knobs();
	init_stats();

	recorder = new AtroposRecorder();

	last_evicted_tracker = NULL;
	last_htentry = NULL;

	/* init learning engine */
	brain_hierarchical = NULL;

    brain_hierarchical = new LearningEngineHierarchical(this,
            knob::atropos_alpha, 
            knob::atropos_gamma, 
            knob::atropos_epsilon,
            knob::atropos_max_actions,
            knob::atropos_seed,
            knob::atropos_policy,
            knob::atropos_learning_type,
            knob::atropos_brain_zero_init);

	bw_level = 0;
	core_ipc = 0;
}

Atropos::~Atropos()
{
	delete brain_hierarchical;
}

void Atropos::print_config()
{
	cout << "atropos_alpha " << knob::atropos_alpha << endl
		<< "atropos_gamma " << knob::atropos_gamma << endl
		<< "atropos_epsilon " << knob::atropos_epsilon << endl
		<< "atropos_state_num_bits " << knob::atropos_state_num_bits << endl
		<< "atropos_max_states " << knob::atropos_max_states << endl
		<< "atropos_seed " << knob::atropos_seed << endl
		<< "atropos_policy " << knob::atropos_policy << endl
		<< "atropos_learning_type " << knob::atropos_learning_type << endl
		<< "atropos_max_actions " << knob::atropos_max_actions << endl
		<< "atropos_pt_size " << knob::atropos_pt_size << endl
		<< "atropos_st_size " << knob::atropos_st_size << endl
		<< "atropos_ht_size " << knob::atropos_ht_size << endl
		<< "atropos_max_pcs " << knob::atropos_max_pcs << endl
		<< "atropos_max_offsets " << knob::atropos_max_offsets << endl
		<< "atropos_max_deltas " << knob::atropos_max_deltas << endl
		<< "atropos_reward_none " << knob::atropos_reward_none << endl
		<< "atropos_reward_incorrect " << knob::atropos_reward_incorrect << endl
		<< "atropos_reward_correct_untimely " << knob::atropos_reward_correct_untimely << endl
		<< "atropos_reward_correct_timely " << knob::atropos_reward_correct_timely << endl
		<< "atropos_reward_expert " << knob::atropos_reward_expert << endl
		<< "atropos_latency_bias" << knob::atropos_latency_bias << endl
		<< "atropos_lookback_window" << knob::atropos_lookback_window << endl
		<< "atropos_brain_zero_init " << knob::atropos_brain_zero_init << endl
		<< "atropos_enable_reward_all " << knob::atropos_enable_reward_all << endl
		<< "atropos_enable_track_multiple " << knob::atropos_enable_track_multiple << endl
		<< "atropos_enable_reward_out_of_bounds " << knob::atropos_enable_reward_out_of_bounds << endl
		<< "atropos_reward_out_of_bounds " << knob::atropos_reward_out_of_bounds << endl
		<< "atropos_state_type " << knob::atropos_state_type << endl
		<< "atropos_state_hash_type " << knob::atropos_state_hash_type << endl
		<< "atropos_access_debug " << knob::atropos_access_debug << endl
		<< "atropos_print_access_debug " << knob::atropos_print_access_debug << endl
		<< "atropos_print_access_debug_pc " << hex << knob::atropos_print_access_debug_pc << dec << endl
		<< "atropos_print_access_debug_pc_count " << knob::atropos_print_access_debug_pc_count << endl
		<< "atropos_print_trace " << knob::atropos_print_trace << endl
		<< "atropos_enable_state_action_stats " << knob::atropos_enable_state_action_stats << endl
		<< "atropos_enable_reward_tracker_hit " << knob::atropos_enable_reward_tracker_hit << endl
		<< "atropos_reward_tracker_hit " << knob::atropos_reward_tracker_hit << endl
		<< "atropos_enable_featurewise_engine " << knob::atropos_enable_featurewise_engine << endl
		<< "atropos_enable_hierarchical_engine " << knob::atropos_enable_hierarchical_engine << endl
		<< "atropos_pref_degree " << knob::atropos_pref_degree << endl
		<< "atropos_enable_dyn_degree " << knob::atropos_enable_dyn_degree << endl
		<< "atropos_max_to_avg_q_thresholds " << array_to_string(knob::atropos_max_to_avg_q_thresholds) << endl
		<< "atropos_dyn_degrees " << array_to_string(knob::atropos_dyn_degrees) << endl
		<< "atropos_multi_deg_select_type " << knob::atropos_multi_deg_select_type << endl
		<< "atropos_last_pref_offset_conf_thresholds " << array_to_string(knob::atropos_last_pref_offset_conf_thresholds) << endl
		<< "atropos_dyn_degrees_type2 " << array_to_string(knob::atropos_dyn_degrees_type2) << endl
		<< "atropos_action_tracker_size " << knob::atropos_action_tracker_size << endl
		<< "atropos_high_bw_thresh " << knob::atropos_high_bw_thresh << endl
		<< "atropos_enable_hbw_reward " << knob::atropos_enable_hbw_reward << endl
		<< "atropos_reward_hbw_correct_timely " << knob::atropos_reward_hbw_correct_timely << endl
		<< "atropos_reward_hbw_correct_untimely " << knob::atropos_reward_hbw_correct_untimely << endl
		<< "atropos_reward_hbw_incorrect " << knob::atropos_reward_hbw_incorrect << endl
		<< "atropos_reward_hbw_none " << knob::atropos_reward_hbw_none << endl
		<< "atropos_reward_hbw_out_of_bounds " << knob::atropos_reward_hbw_out_of_bounds << endl
		<< "atropos_reward_hbw_tracker_hit " << knob::atropos_reward_hbw_tracker_hit << endl
		<< "atropos_last_pref_offset_conf_thresholds_hbw " << array_to_string(knob::atropos_last_pref_offset_conf_thresholds_hbw) << endl
		<< "atropos_dyn_degrees_type2_hbw " << array_to_string(knob::atropos_dyn_degrees_type2_hbw) << endl
		<< endl
		<< "le_enable_trace " << knob::le_enable_trace << endl
		<< "le_trace_interval " << knob::le_trace_interval << endl
		<< "le_trace_file_name " << knob::le_trace_file_name << endl
		<< "le_trace_state " << hex << knob::le_trace_state << dec << endl
		<< "le_enable_score_plot " << knob::le_enable_score_plot << endl
		<< "le_plot_file_name " << knob::le_plot_file_name << endl
		<< "le_plot_actions " << array_to_string(knob::le_plot_actions) << endl
		<< "le_enable_action_trace " << knob::le_enable_action_trace << endl
		<< "le_action_trace_interval " << knob::le_action_trace_interval << endl
		<< "le_action_trace_name " << knob::le_action_trace_name << endl
		<< "le_enable_action_plot " << knob::le_enable_action_plot << endl
		<< endl
		<< "le_featurewise_active_features " << print_active_features2(knob::le_featurewise_active_features) << endl
		<< "le_featurewise_active_worker " << print_active_features2(knob::le_featurewise_active_worker) << endl
		<< "le_featurewise_num_tilings " << array_to_string(knob::le_featurewise_num_tilings) << endl
		<< "le_featurewise_num_tiles " << array_to_string(knob::le_featurewise_num_tiles) << endl
		<< "le_featurewise_hash_types " << array_to_string(knob::le_featurewise_hash_types) << endl
		<< "le_featurewise_enable_tiling_offset " << array_to_string(knob::le_featurewise_enable_tiling_offset) << endl
		<< "le_featurewise_max_q_thresh " << knob::le_featurewise_max_q_thresh << endl
		<< "le_featurewise_enable_action_fallback " << knob::le_featurewise_enable_action_fallback << endl
		<< "le_featurewise_feature_weights " << array_to_string(knob::le_featurewise_feature_weights) << endl
		<< "le_featurewise_enable_dynamic_weight " << knob::le_featurewise_enable_dynamic_weight << endl
		<< "le_featurewise_weight_gradient " << knob::le_featurewise_weight_gradient << endl
		<< "le_featurewise_disable_adjust_weight_all_features_align " << knob::le_featurewise_disable_adjust_weight_all_features_align << endl
		<< "le_featurewise_selective_update " << knob::le_featurewise_selective_update << endl
		<< "le_featurewise_pooling_type " << knob::le_featurewise_pooling_type << endl
		<< "le_featurewise_enable_dyn_action_fallback " << knob::le_featurewise_enable_dyn_action_fallback << endl
		<< "le_featurewise_bw_acc_check_level " << knob::le_featurewise_bw_acc_check_level << endl
		<< "le_featurewise_acc_thresh " << knob::le_featurewise_acc_thresh << endl
		<< "le_featurewise_enable_trace " << knob::le_featurewise_enable_trace << endl
		<< "le_featurewise_trace_feature_type " << knob::le_featurewise_trace_feature_type << endl
		<< "le_featurewise_trace_feature " << knob::le_featurewise_trace_feature << endl
		<< "le_featurewise_trace_interval " << knob::le_featurewise_trace_interval << endl
		<< "le_featurewise_trace_record_count " << knob::le_featurewise_trace_record_count << endl
		<< "le_featurewise_trace_file_name " << knob::le_featurewise_trace_file_name << endl
		<< "le_featurewise_enable_score_plot " << knob::le_featurewise_enable_score_plot << endl
		<< "le_featurewise_plot_actions " << array_to_string(knob::le_featurewise_plot_actions) << endl
		<< "le_featurewise_plot_file_name " << knob::le_featurewise_plot_file_name << endl
		<< "le_page_shifting " << knob::le_page_shifting << endl
		<< "le_offset_shifting " << knob::le_featurewise_plot_file_name << endl
		<< endl;
}

void Atropos::invoke_prefetcher(uint64_t pc, uint64_t address, uint8_t cache_hit, uint8_t type, vector<uint64_t> &pref_addr)
{
	uint64_t page = address >> LOG2_PAGE_SIZE;
	uint32_t offset = (address >> LOG2_BLOCK_SIZE) & ((1ull << (LOG2_PAGE_SIZE - LOG2_BLOCK_SIZE)) - 1);

	MYLOG("---------------------------------------------------------------------");
	MYLOG("%s %lx pc %lx page %lx off %u", GetAccessType(type), address, pc, page, offset);

	/* compute reward on demand */
	reward(address);

	/* record the access: just to gain some insights from the workload
	 * defined in atropos_helper.h(cc) */
	recorder->record_access(pc, address, page, offset, bw_level);

	/* global state tracking */
	Atropos_STEntry *stentry = update_global_state(pc, page, offset, address);
	/* per page state tracking */
	// Atropos_STEntry *stentry = update_local_state(pc, page, offset, address);

	/* Measure state.
	 * state can contain per page local information like delta signature, pc signature etc.
	 * it can also contain global signatures like last three branch PCs etc.
	 */
	State *state = new State();
	state->pc = pc;
	state->address = address;
	state->page = page;
	state->offset = offset;
	state->delta = !stentry->deltas.empty() ? stentry->deltas.back() : 0;
	state->local_delta_sig = stentry->get_delta_sig();
	state->local_delta_sig2 = stentry->get_delta_sig2();
	state->local_pc_sig = stentry->get_pc_sig();
	state->local_offset_sig = stentry->get_offset_sig();
	state->bw_level = bw_level;
	state->is_high_bw = is_high_bw();
	state->acc_level = acc_level;

	uint32_t count = pref_addr.size();
	predict(address, page, offset, state, pref_addr);
    history_track(address, state);
	stats.pref_issue.atropos += (pref_addr.size() - count);
}

Atropos_STEntry* Atropos::update_global_state(uint64_t pc, uint64_t page, uint32_t offset, uint64_t address)
{
	stats.st.lookup++;
	Atropos_STEntry *stentry = NULL;
	auto st_index = signature_table.begin(); // 任意page
	if(st_index != signature_table.end())
	{
		stats.st.hit++;
		stentry = (*st_index);
		stentry->update(page, pc, offset, address);
		signature_table.erase(st_index);
		signature_table.push_back(stentry);
		return stentry;
	}
	else
	{
		stats.st.insert++;
		stentry = new Atropos_STEntry(page, pc, offset);
		recorder->record_trigger_access(page, pc, offset);
		signature_table.push_back(stentry);
		return stentry;
	}
}


Atropos_STEntry* Atropos::update_local_state(uint64_t pc, uint64_t page, uint32_t offset, uint64_t address)
{
	stats.st.lookup++;
	Atropos_STEntry *stentry = NULL;
	auto st_index = find_if(signature_table.begin(), signature_table.end(), [page](Atropos_STEntry *stentry){return stentry->page == page;});
	if(st_index != signature_table.end())
	{
		stats.st.hit++;
		stentry = (*st_index);
		stentry->update(page, pc, offset, address);
		signature_table.erase(st_index);
		signature_table.push_back(stentry);
		return stentry;
	}
	else
	{
		if(signature_table.size() >= knob::atropos_st_size)
		{
			stats.st.evict++;
			stentry = signature_table.front();
			signature_table.pop_front();
			if(knob::atropos_access_debug)
			{
				recorder->record_access_knowledge(stentry);
				if(knob::atropos_print_access_debug)
				{
					// print_access_debug(stentry);
				}
			}
			delete stentry;
		}

		stats.st.insert++;
		stentry = new Atropos_STEntry(page, pc, offset);
		recorder->record_trigger_access(page, pc, offset);
		signature_table.push_back(stentry);
		return stentry;
	}
}

uint32_t Atropos::predict(uint64_t base_address, uint64_t page, uint32_t offset, State *state, vector<uint64_t> &pref_addr)
{
	MYLOG("addr@%lx page %lx off %u state %x", base_address, page, offset, state->value());

	stats.predict.called++;

	/* query learning engine to get the next prediction */
    uint32_t page_option = brain_hierarchical->choosePageOption(state, 0); // 0-7
    uint32_t page_action = brain_hierarchical->choosePageAction(state, page_option); // 0-63

    uint32_t offset_option = brain_hierarchical->chooseOffsetOption(state, 0); // 0-7
    uint32_t offset_action = brain_hierarchical->chooseOffsetAction(state, offset_option); // 0-63

    // PGOPT: page_option
    // PGACT: page_action
    // OFOPT: offset_option
    // OFACT: offset_action
    MYLOG("worker action taken PGOPT: %u, PGACT: %u, OFOPT: %u, OFACT: %u, STATE: %s, SCORES: %s", 
        page_option, page_action, offset_option, offset_action, state->to_string().c_str(), getStringQ(state).c_str());

	// MYLOG("[PREDICT] PAGE: %u, OFFSET: %d", PAGE_DELTA[page_action], OFFSET_DELTA[offset_action]);

	uint64_t addr = 0xdeadbeef;
	Atropos_PTEntry *ptentry = NULL;
	int32_t predicted_page = 0;
	int32_t predicted_offset = 0;

	if (page_action != 0 || offset_action != 0) {
		predicted_page = (int32_t)page + (page_action - ACTION_SHIFT);
		predicted_offset = action2offset(offset_action, offset);

		if (predicted_page >= 0) {
            /* get predicted addr */
			addr = (predicted_page << LOG2_PAGE_SIZE) + (predicted_offset << LOG2_BLOCK_SIZE);
			MYLOG("[PRED] PAGE: %d, OFFSET: %d, ADDR: %lx", predicted_page, predicted_offset, addr);

			/* track prefetch */
            // @zhangtao: 注意，track中，evict PTEntry后，也会调用train
			bool new_addr = track(addr, state, page_option, page_action, offset_option, offset_action, &ptentry);
			if (new_addr)
			{
				pref_addr.push_back(addr);
			}
		} else {
			MYLOG("pred_off %d out_of_bounds", predicted_offset);
			stats.predict.out_of_bounds++;
	        track(addr, state, page_option, page_action, offset_option, offset_action, &ptentry);
		}
	} else {
		MYLOG("no prefecth");
		/* agent decided not to prefetch */
		addr = 0xdeadbeef;
		/* track no prefetch */
        // @zhangtao 不预取为什么还要track
	    // track(addr, state, page_option, page_action, offset_option, offset_action, &ptentry);
	}

	stats.predict.predicted += pref_addr.size();
	MYLOG("end@%lx", base_address);

	return pref_addr.size();
}

/* Returns true if the address is not already present in prefetch_tracker
 * false otherwise */
bool Atropos::track(uint64_t address, State *state, uint32_t page_option, uint32_t page_action, 
    uint32_t offset_option, uint32_t offset_action, Atropos_PTEntry **tracker)
{
	// MYLOG("addr@%lx state %x act_idx %u act %d", address, state->value(), action_index, Actions[action_index]);
	stats.track.called++;

	bool new_addr = true;
	vector<Atropos_PTEntry*> ptentries = search_pt(address, false);
	if(ptentries.empty())
		new_addr = true;
	else
		new_addr = false;

	if(!new_addr && address != 0xdeadbeef && !knob::atropos_enable_track_multiple)
	{
		stats.track.same_address++;
		tracker = NULL;
		return new_addr;
	}

	/* new prefetched address that hasn't been seen before */
	Atropos_PTEntry *ptentry = NULL;

	if(prefetch_tracker.size() >= knob::atropos_pt_size)
	{
		stats.track.evict++;
		ptentry = prefetch_tracker.front();
		prefetch_tracker.pop_front();
		// MYLOG("victim_state %x victim_act_idx %u victim_act %d", ptentry->state->value(), ptentry->action_index, Actions[ptentry->action_index]);
		if(last_evicted_tracker)
		{
			// MYLOG("last_victim_state %x last_victim_act_idx %u last_victim_act %d", last_evicted_tracker->state->value(), last_evicted_tracker->action_index, Actions[last_evicted_tracker->action_index]);
			/* train the agent */
			train(ptentry, last_evicted_tracker);
			delete last_evicted_tracker->state;
			delete last_evicted_tracker;
		}
		last_evicted_tracker = ptentry;
	}

	ptentry = new Atropos_PTEntry(address, state, page_option, page_action, offset_option, offset_action);
	prefetch_tracker.push_back(ptentry);
	assert(prefetch_tracker.size() <= knob::atropos_pt_size);

	(*tracker) = ptentry;
	MYLOG("end@%lx", address);

	return new_addr;
}


void Atropos::history_track(uint64_t address, State *state)
{
	// MYLOG("addr@%lx state %x act_idx %u act %d", address, state->value(), action_index, Actions[action_index]);

	Atropos_HTEntry *htentry = NULL;

	if(history_table.size() >= knob::atropos_ht_size)
	{
		htentry = history_table.front();
		history_table.pop_front();
		if(htentry)
		{
			delete htentry->state;
			delete htentry;
		}
		htentry = NULL;
	}

	htentry = new Atropos_HTEntry(address, state);
	history_table.push_back(htentry);

	MYLOG("history_track @%lx", address);

	return;
}


/* This reward fucntion is called after seeing a demand access to the address */
/* TODO: what if multiple prefetch request generated the same address?
 * Currently, it just rewards the oldest prefetch request to the address
 * Should we reward all? */
void Atropos::reward(uint64_t address)
{
	MYLOG("addr @ %lx", address);

	stats.reward.demand.called++;
	vector<Atropos_PTEntry*> ptentries = search_pt(address, knob::atropos_enable_reward_all);

	if(ptentries.empty())
	{
		MYLOG("PT miss");
		stats.reward.demand.pt_not_found++;
		return;
	}
	else
	{
		stats.reward.demand.pt_found++;
	}

	for(uint32_t index = 0; index < ptentries.size(); ++index)
	{
		Atropos_PTEntry *ptentry = ptentries[index];
		stats.reward.demand.pt_found_total++;

		// MYLOG("PT hit. state %x act_idx %u act %d", ptentry->state->value(), ptentry->action_index, Actions[ptentry->action_index]);
		/* Do not compute reward if already has a reward.
		 * This can happen when a prefetch access sees multiple demand reuse */
		if(ptentry->has_reward)
		{
			MYLOG("entry already has reward: %d", ptentry->reward);
			stats.reward.demand.has_reward++;
			return;
		}

		if(ptentry->is_filled) /* timely */
		{
			assign_reward(ptentry, RewardType::correct_timely);
			MYLOG("assigned reward correct_timely(%d)", ptentry->reward);
		}
		else
		{
			assign_reward(ptentry, RewardType::correct_untimely);
			MYLOG("assigned reward correct_untimely(%d)", ptentry->reward);
		}
		ptentry->has_reward = true;
	}
}

/* This reward function is called during eviction from prefetch_tracker */
void Atropos::reward(Atropos_PTEntry *ptentry)
{
	// MYLOG("reward PT evict %lx state %x act_idx %u act %d", ptentry->address, ptentry->state->value(), ptentry->action_index, Actions[ptentry->action_index]);

	stats.reward.train.called++;
	assert(!ptentry->has_reward);
	/* this is called during eviction from prefetch tracker
	 * that means, this address doesn't see a demand reuse.
	 * hence it either can be incorrect, or no prefetch */
	if(ptentry->address == 0xdeadbeef) /* no prefetch */
	{
		assign_reward(ptentry, RewardType::none);
		MYLOG("assigned reward no_pref(%d)", ptentry->reward);
	}
	else /* incorrect prefetch */
	{
		assign_reward(ptentry, RewardType::incorrect);
		MYLOG("assigned reward incorrect(%d)", ptentry->reward);
	}
	ptentry->has_reward = true;
}

void Atropos::assign_reward(Atropos_PTEntry *ptentry, RewardType type)
{
	// MYLOG("assign_reward PT evict %lx state %x act_idx %u act %d", ptentry->address, ptentry->state->value(), ptentry->action_index, Actions[ptentry->action_index]);
	assert(!ptentry->has_reward);

	/* compute the reward */
	int32_t reward = compute_reward(ptentry, type);

	/* assign */
	ptentry->reward = reward;
	ptentry->reward_type = type;
	ptentry->has_reward = true;

	/* maintain stats */
	stats.reward.assign_reward.called++;
	switch(type)
	{
		case RewardType::correct_timely: 	stats.reward.correct_timely++; break;
		case RewardType::correct_untimely: 	stats.reward.correct_untimely++; break;
		case RewardType::incorrect: 		stats.reward.incorrect++; break;
		case RewardType::none: 				stats.reward.no_pref++; break;
		case RewardType::out_of_bounds: 	stats.reward.out_of_bounds++; break;
		case RewardType::tracker_hit: 		stats.reward.tracker_hit++; break;
		default:							assert(false);
	}
}

int32_t Atropos::compute_reward(Atropos_PTEntry *ptentry, RewardType type)
{
	bool high_bw = (knob::atropos_enable_hbw_reward && is_high_bw()) ? true : false;
	int32_t reward = 0;

	stats.reward.compute_reward.dist[type][high_bw]++;

	if(type == RewardType::correct_timely)
	{
		reward = high_bw ? knob::atropos_reward_hbw_correct_timely : knob::atropos_reward_correct_timely;
	}
	else if(type == RewardType::correct_untimely)
	{
		reward = high_bw ? knob::atropos_reward_hbw_correct_untimely : knob::atropos_reward_correct_untimely;
	}
	else if(type == RewardType::incorrect)
	{
		reward = high_bw ? knob::atropos_reward_hbw_incorrect : knob::atropos_reward_incorrect;
	}
	else if(type == RewardType::none)
	{
		reward = high_bw ? knob::atropos_reward_hbw_none : knob::atropos_reward_none;
	}
	else if(type == RewardType::out_of_bounds)
	{
		reward = high_bw ? knob::atropos_reward_hbw_out_of_bounds : knob::atropos_reward_out_of_bounds;
	}
	else if(type == RewardType::tracker_hit)
	{
		reward = high_bw ? knob::atropos_reward_hbw_tracker_hit : knob::atropos_reward_tracker_hit;
	}
	else
	{
		cout << "Invalid reward type found " << type << endl;
		assert(false);
	}

	return reward;
}

// behaviour clone
void Atropos::train_behaviour_clone(Atropos_HTEntry *curr_evicted, Atropos_HTEntry *last_evicted)
{
	// MYLOG("train_bc victim %s %u %d last_victim %s %u %d", curr_evicted->state->to_string().c_str(), curr_evicted->action_index, Actions[curr_evicted->action_index], last_evicted->state->to_string().c_str(), last_evicted->action_index, Actions[last_evicted->action_index]);

    brain_hierarchical->learn(last_evicted->state, 
            last_evicted->action, 
            last_evicted->reward, 
            curr_evicted->state, 
            curr_evicted->action, 
            last_evicted->reward_type);
}


void Atropos::train(Atropos_PTEntry *curr_evicted, Atropos_PTEntry *last_evicted)
{
	// MYLOG("victim %s %u %d last_victim %s %u %d", curr_evicted->state->to_string().c_str(), curr_evicted->action_index, Actions[curr_evicted->action_index], last_evicted->state->to_string().c_str(), last_evicted->action_index, Actions[last_evicted->action_index]);

	stats.train.called++;
	if(!last_evicted->has_reward)
	{
		stats.train.compute_reward++;
		reward(last_evicted);
	}
	assert(last_evicted->has_reward);

	/* train */
	// MYLOG("===QLEARNING=== S1: %s A1: %u R1: %d S2: %s A2: %u", last_evicted->state->to_string().c_str(), last_evicted->action_index, last_evicted->reward, curr_evicted->state->to_string().c_str(), curr_evicted->action_index);

    brain_hierarchical->learn(last_evicted->state, 
            last_evicted->action, 
            last_evicted->reward, 
            curr_evicted->state, 
            curr_evicted->action, 
            last_evicted->reward_type);

	MYLOG("train done");
}

/* TODO: what if multiple prefetch request generated the same address?
 * Currently it just sets the fill bit of the oldest prefetch request.
 * Do we need to set it for everyone? */
void Atropos::register_fill(uint64_t address)
{
	MYLOG("fill @ %lx", address);

	stats.register_fill.called++;
	vector<Atropos_PTEntry*> ptentries = search_pt(address, knob::atropos_enable_reward_all);
	if(!ptentries.empty())
	{
		stats.register_fill.set++;
		for(uint32_t index = 0; index < ptentries.size(); ++index)
		{
			stats.register_fill.set_total++;
			ptentries[index]->is_filled = true;
			// MYLOG("fill PT hit. pref with act_idx %u act %d", ptentries[index]->action_index, Actions[ptentries[index]->action_index]);
		}
	}
}

// DELTA_SHIFT = 32
// action = delta + DELTA_SHIFT // [0,63]
// delta = action - DELTA_SHIFT // [-32,31]
// offset = offset' + delta

// [-63,63] -> [-32,31] -> [0,63]
int Atropos::delta2action(int delta)
{
    if (delta > 31) { // e.g. 40
        delta = delta - 64; // e.g. -24 // [-63,63] -> [-32,31]
    } else if (delta < -32) { // e.g. -40
        delta = delta + 64; // e.g. 24  // [-63,63] -> [-32,31]
    }

    return delta + ACTION_SHIFT; // delta + 32 // e.g. -24 + 32 = 8 // [0,63]
}



// [0,63] -> [-32,31] -> [-126, 126] -> [0,63]
int Atropos::action2offset(int action, int offset)
{
    int delta = action - ACTION_SHIFT; // [0,63] -> [-32,31]
    offset = offset + delta;
    if (offset > 63) { // e.g. 80
        offset = offset - 64; // e.g. 80 - 64 = 16
    } else if (offset < 0) { // e.g. -40
        offset = offset + 64; // e.g. 24  // [-32,31]
    }
    return offset;
}


// 在cache fill时被调用，
void Atropos::update_htentry(uint64_t address, uint64_t latency)
{
	MYLOG("update_htentry @ %lx", address);

	Atropos_HTEntry *htentry = NULL;
    Atropos_HTEntry *entangle_htentry = NULL;
    int index = search_ht(address, htentry);
    int page_delta = 65535;

	if (htentry != NULL) {
        htentry->latency = latency;

        int begin_index = index - latency / (L1D_LATENCY + L2C_LATENCY) - 1 + knob::atropos_latency_bias;
        if (begin_index < 0) return;

        int end_index = begin_index - knob::atropos_lookback_window;
        int step_index = begin_index;

        while (step_index > 0 && step_index > end_index) {

            Atropos_HTEntry *tmp_htentry = get_htentry(step_index);
            int page_delta = (htentry->address >> LOG2_PAGE_SIZE) - (tmp_htentry->address >> LOG2_PAGE_SIZE);

            if (page_delta <= 31 && page_delta >= -32) break;

            step_index -= 1;
        }

        if (page_delta > 31 || page_delta < -32) return;

        entangle_htentry = get_htentry(step_index);

        int offset_delta = ((htentry->address - entangle_htentry->address) >> LOG2_BLOCK_SIZE) & ((1ull << (LOG2_PAGE_SIZE - LOG2_BLOCK_SIZE)) - 1); // [-63, 63]

        int page_action = page_delta + ACTION_SHIFT; // [0,63]
        int offset_action = delta2action(offset_delta); // [0,63]

        htentry->action.page_option = page_action / 8;
        htentry->action.page_action = page_action;

        htentry->action.offset_option = offset_action / 8;
        htentry->action.offset_action = offset_action;

        htentry->reward = knob::atropos_reward_expert;
        htentry->reward_type = RewardType::expert;
        htentry->has_reward = true;

        if (last_htentry != NULL) {
            train_behaviour_clone(htentry, last_htentry); // 只用第一项
            delete last_htentry->state;
            delete last_htentry;
        }

        last_htentry = htentry;
	}
}



void Atropos::register_prefetch_hit(uint64_t address)
{
	MYLOG("pref_hit @ %lx", address);

	stats.register_prefetch_hit.called++;
	vector<Atropos_PTEntry*> ptentries = search_pt(address, knob::atropos_enable_reward_all);
	if(!ptentries.empty())
	{
		stats.register_prefetch_hit.set++;
		for(uint32_t index = 0; index < ptentries.size(); ++index)
		{
			stats.register_prefetch_hit.set_total++;
			ptentries[index]->pf_cache_hit = true;
			// MYLOG("pref_hit PT hit. pref with act_idx %u act %d", ptentries[index]->action_index, Actions[ptentries[index]->action_index]);
		}
	}
}

vector<Atropos_PTEntry*> Atropos::search_pt(uint64_t address, bool search_all)
{
	vector<Atropos_PTEntry*> entries;
	for(uint32_t index = 0; index < prefetch_tracker.size(); ++index)
	{
		if(prefetch_tracker[index]->address == address)
		{
			entries.push_back(prefetch_tracker[index]);
			if(!search_all) break;
		}
	}
	return entries;
}

vector<Atropos_HTEntry*> Atropos::search_ht(uint64_t address, bool search_all)
{
	vector<Atropos_HTEntry*> entries;
	for(uint32_t index = 0; index < history_table.size(); ++index)
	{
		if(history_table[index]->address == address)
		{
			entries.push_back(history_table[index]);
			if(!search_all) break;
		}
	}
	return entries;
}

int Atropos::search_ht(uint64_t address, Atropos_HTEntry* htentry)
{
	for(uint32_t index = 0; index < history_table.size(); ++index)
	{
		if(history_table[index]->address == address)
		{
            htentry = history_table[index];
            return index;
		}
	}
    return -1; // NOT FOUND
}

Atropos_HTEntry *Atropos::get_htentry(int index)
{
    if (index < 0 || index >= history_table.size()) return NULL;
    return history_table[index];
}

void Atropos::update_bw(uint8_t bw)
{
	assert(bw < DRAM_BW_LEVELS);
	bw_level = bw;
	stats.bandwidth.epochs++;
	stats.bandwidth.histogram[bw_level]++;
}

void Atropos::update_ipc(uint8_t ipc)
{
	assert(ipc < ATROPOS_MAX_IPC_LEVEL);
	core_ipc = ipc;
	stats.ipc.epochs++;
	stats.ipc.histogram[ipc]++;
}

void Atropos::update_acc(uint32_t acc)
{
	assert(acc < CACHE_ACC_LEVELS);
	acc_level = acc;
	stats.cache_acc.epochs++;
	stats.cache_acc.histogram[acc]++;
}

bool Atropos::is_high_bw()
{
	return bw_level >= knob::atropos_high_bw_thresh ? true : false;
}

void Atropos::dump_stats()
{

}
